package kafka

import (
	"context"
	"encoding/json"
	"fmt"

	"github.com/example-user/go-pocket-etl/pkg/core"
	"github.com/example-user/go-pocket-etl/pkg/factory"
	"github.com/segmentio/kafka-go"
	"github.com/sirupsen/logrus"
)

// init 函数在包被加载时自动执行，用于向工厂注册此 Sink 组件。
func init() {
	factory.RegisterSink("kafka", NewSink)
}

// Sink 实现了 core.Sink 接口，用于将数据记录作为消息写入到 Kafka topic。
type Sink struct {
	writer *kafka.Writer // 来自 segmentio/kafka-go 库的 writer 实例
}

// NewSink 是 kafka.Sink 的构造函数。
func NewSink() core.Sink {
	return &Sink{}
}

// Open 负责初始化 Kafka writer。
// 它需要从配置中读取 'brokers' (一个地址列表) 和 'topic'。
func (s *Sink) Open(config map[string]interface{}) error {
	brokersVal, ok := config["brokers"]
	if !ok {
		return fmt.Errorf("kafka sink: config is missing 'brokers'")
	}
	brokers, ok := brokersVal.([]interface{})
	if !ok || len(brokers) == 0 {
		return fmt.Errorf("kafka sink: 'brokers' must be a non-empty array of strings")
	}

	brokerAddresses := make([]string, len(brokers))
	for i, b := range brokers {
		brokerAddresses[i], ok = b.(string)
		if !ok {
			return fmt.Errorf("kafka sink: 'brokers' array must contain only strings")
		}
	}

	topic, ok := config["topic"].(string)
	if !ok || topic == "" {
		return fmt.Errorf("kafka sink: config is missing or has invalid 'topic'")
	}

	// 配置 kafka.Writer。
	// kafka-go 的 Writer 内部已经实现了连接管理、自动批处理和重试逻辑。
	s.writer = &kafka.Writer{
		Addr:     kafka.TCP(brokerAddresses...),
		Topic:    topic,
		Balancer: &kafka.LeastBytes{}, // 使用最小字节分发器，将消息发送到分区负载最低的 broker
		Async:    false, // 设置为 false (同步模式) 以确保 WriteMessages 是阻塞的，这使得错误处理更直接
	}

	logrus.Infof("Kafka sink initialized for topic '%s' with brokers %v", topic, brokerAddresses)
	return nil
}

// Write 将一批记录序列化为 JSON，并作为消息批量写入到 Kafka。
func (s *Sink) Write(records []core.Record) error {
	if s.writer == nil {
		return fmt.Errorf("kafka sink: writer is not initialized, was Open() called?")
	}
	if len(records) == 0 {
		return nil
	}

	messages := make([]kafka.Message, 0, len(records))

	for _, record := range records {
		// 将整个 record 序列化为 JSON 字节，作为 Kafka 消息的 Value。
		jsonValue, err := json.Marshal(record)
		if err != nil {
			// 记录序列化错误，但选择跳过这条有问题的记录，继续处理批次中的其他记录。
			// 这是一种更具弹性的错误处理策略。
			logrus.Errorf("Failed to marshal record to JSON, skipping: %v, record: %v", err, record)
			continue
		}

		messages = append(messages, kafka.Message{
			// Key 可以用于 Kafka 的分区策略，确保具有相同 Key 的消息进入同一个分区。
			// 在这个通用实现中，我们暂时不设置 Key。
			Value: jsonValue,
		})
	}

	// 如果所有记录都因为序列化失败而被跳过，则无需执行写入操作。
	if len(messages) == 0 {
		return nil
	}

	// 使用上下文进行写入操作，以便未来可以支持超时控制。
	ctx := context.Background()

	err := s.writer.WriteMessages(ctx, messages...)
	if err != nil {
		return fmt.Errorf("kafka sink: failed to write messages: %w", err)
	}

	logrus.Debugf("Successfully wrote %d messages to Kafka", len(messages))
	return nil
}

// Close 关闭 Kafka writer，这个操作会阻塞直到所有缓冲区的消息都被成功发送。
func (s *Sink) Close() error {
	if s.writer != nil {
		logrus.Info("Closing Kafka writer...")
		return s.writer.Close()
	}
	return nil
}
